#!/usr/bin/env python

import configparser
import fnmatch
import os
from argparse import ArgumentParser
import numpy as np
from astropy.io import fits

"""

"""


def _get_value(config, name):
    value = config.get('FITS_HEADER', name)
    if name.upper() in ['H_TRANS', 'V_TRANS', 'IMGINDEX', 'SEQNUM']:
        value = int(value)
    if name.upper() in ['DARKTIME', 'EXPTIME']:
        value = float(value)
    if name.upper() in ['WAVELEN', 'SYSTEMP']:
        value = np.float32(value)
    return value


def readinfo(finfo):
    output = dict()
    cf = configparser.ConfigParser()
    cf.read(finfo)
    if cf.has_section('FITS_HEADER') is False:
        return output
    options = cf.options('FITS_HEADER')
    for opt in options:
        output[opt] = _get_value(cf, opt)
    return output


def bin2fits(binfname, outfits=None, overwrite=True, \
             nx=1250, ny=4700, dtype=np.uint16, header_info=None):
    """
    TODO:
    1) 增加头文件内容
    2）修改为直接将预存的数组转换为fits
    3) 增加fits输出路径
    """
#    print('*** bin2fits ***')

    # fout = binfname.replace('.{}'.format(suffix), '.fits')
    fout = outfits

    output = np.zeros((ny * 2, nx * 8), dtype=dtype)
    data = np.fromfile(binfname, dtype=dtype)
    data = data.astype(np.uint16)
    data = data.reshape((ny * nx, 16))

    for chan in range(16):
        y = chan // 8
        if y == 0:
            x = chan
        else:
            x = 15 - chan
        x0 = x * nx
        y0 = y * ny

        dd = data[:, chan].reshape((ny, nx))
        if y == 1:
            dd = dd[::-1, :]
        if x > 3:
            dd = dd[:, ::-1]
        output[y0:y0+ny, x0:x0+nx] = dd

    hdu = fits.PrimaryHDU(output)

    # add header info
    if header_info is not None:
        for key in header_info:
#            print('@@@@@@ key: {}, val: {}'.format(key, header_info[key]))
            hdu.header[key.upper()] = header_info[key]

    hdu.writeto(fout, overwrite=overwrite)
#    print('write fits file: {}'.format(fout))


def main(path, **kwargs):
    if os.path.isfile(path):
        bin2fits(path, **kwargs)

    elif os.path.isdir(path):
        if path.endswith('/'):
            path = path[:-1]
        pattern = '*.{}'.format(kwargs['suffix'])
        flist = []
        for r, d, f in os.walk(path):
            names = fnmatch.filter(f, pattern)
            if len(names) > 0:
                for name in names:
                    flist.append('{}/{}'.format(r, name))
        for f in flist:
            bin2fits(f, **kwargs)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('path', type=str,
                        help='file name or directory name (recursively search binary files)')
    parser.add_argument('--overwrite', '-o', action='store_true',
                        help='if set, overwrite existing files')
    parser.add_argument('--suffix', '-s', default='bin', type=str,
                        help='suffix of binary files, default is "bin"')
    parser.add_argument('--col', '-x', default=1250, type=int,
                        help='number of columns in each channel, default is 1250')
    parser.add_argument('--row', '-y', default=4700, type=int,
                        help='number of rows in each channel, default is 4700')
    parser.add_argument('--dtype', '-d', default='float32', type=str,
                        help='binary data type, default is float32')
    args = parser.parse_args()
    main(args.path, overwrite=args.overwrite, suffix=args.suffix, nx=args.col, ny=args.row, dtype=args.dtype)
